import os
import tushare as ts
import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error
from tensorflow.keras.models import Sequential, load_model
from tensorflow.keras.layers import SimpleRNN, Dense, Dropout, Input
import matplotlib
import matplotlib.pyplot as plt
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 支持中文
matplotlib.rcParams['axes.unicode_minus'] = False    # 正常显示负号
from PIL import Image
import gradio as gr
import io

# 创建保存模型的文件夹
MODEL_DIR = "saved_models"
os.makedirs(MODEL_DIR, exist_ok=True)

# 获取股票数据的函数（前复权）
pro = ts.pro_api('c396c2474256e114c94734e2c48bdc9073313c2332e63375fd7845f9')

def get_stock_data(stock_code, start_date='20000101', end_date='20250501'):
    df = pro.daily(
        ts_code=stock_code,
        start_date=start_date,
        end_date=end_date
    )
    if df.empty:
        raise ValueError(f"{stock_code} 没有获取到数据")
    df = df[['trade_date', 'close']]
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df.set_index('trade_date', inplace=True)
    df = df.sort_index()
    return df

# 数据预处理：归一化
def preprocess_data(df):
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(df['close'].values.reshape(-1, 1))
    return scaled_data, scaler

# 创建时间序列数据集
def create_dataset(data, time_step):
    X, y = [], []
    for i in range(time_step, len(data)):
        X.append(data[i - time_step:i, 0])
        y.append(data[i, 0])
    return np.array(X), np.array(y)

# 构建 RNN 模型
def build_rnn_model(input_shape):
    model = Sequential()
    model.add(Input(shape=input_shape))
    model.add(SimpleRNN(units=50, return_sequences=True))
    model.add(Dropout(0.2))
    model.add(SimpleRNN(units=50, return_sequences=False))
    model.add(Dropout(0.2))
    model.add(Dense(units=1))
    model.compile(optimizer='adam', loss='mean_squared_error')
    return model

# 训练并预测函数
def train_and_predict(stock_code, time_range='全部', start_date=None, end_date=None, forecast_days=30):
    # 获取股票数据
    if start_date and end_date:
        stock_data = get_stock_data(stock_code, start_date, end_date)
    else:
        stock_data = get_stock_data(stock_code)
    # 根据时间范围筛选数据
    if time_range not in ['全部', '自定义输入']:
        years = {'一年': 1, '三年': 3, '五年': 5, '十年': 10}[time_range]
        start_date = stock_data.index.max() - pd.DateOffset(years=years)
        end_date = stock_data.index.max()
        stock_data = stock_data[stock_data.index >= start_date]
    elif time_range == '全部':
        start_date = stock_data.index.min()
        end_date = stock_data.index.max()
    # 统一格式化日期
    start_str = pd.to_datetime(start_date).strftime('%Y%m%d')
    end_str = pd.to_datetime(end_date).strftime('%Y%m%d')
    model_path = os.path.join(MODEL_DIR, f"{stock_code}_{start_str}_{end_str}_model.h5")

    # 数据预处理
    scaled_data, scaler = preprocess_data(stock_data)

    # 创建时间序列数据集
    time_step = 120  # 使用过去120天的数据来预测
    X, y = create_dataset(scaled_data, time_step)

    # 划分训练集和测试集
    train_size = int(len(X) * 0.8)
    X_train, X_test = X[:train_size], X[train_size:]
    y_train, y_test = y[:train_size], y[train_size:]

    # Reshape 输入数据为 RNN 格式 (样本数, 时间步长, 特征数)
    X_train = X_train.reshape(X_train.shape[0], X_train.shape[1], 1)
    X_test = X_test.reshape(X_test.shape[0], X_test.shape[1], 1)

    # 检查是否已有保存的模型
    if os.path.exists(model_path):
        model = load_model(model_path)
    else:
        model = build_rnn_model((X_train.shape[1], 1))
        model.fit(X_train, y_train, epochs=10, batch_size=32, verbose=0)
        model.save(model_path)

    # 进行预测
    predicted_stock_price = model.predict(X_test)
    predicted_stock_price = scaler.inverse_transform(predicted_stock_price)
    real_stock_price = scaler.inverse_transform(y_test.reshape(-1, 1))

    # 计算均方误差
    mse = mean_squared_error(real_stock_price, predicted_stock_price)

    # 未来股价预测
    last_seq = scaled_data[-time_step:].reshape(1, time_step, 1)
    future_predictions = []
    for _ in range(forecast_days):
        next_pred = model.predict(last_seq)
        future_predictions.append(next_pred[0, 0])
        last_seq = np.append(last_seq[:, 1:, :], next_pred.reshape(1, 1, 1), axis=1)
    future_predictions = scaler.inverse_transform(np.array(future_predictions).reshape(-1, 1)).flatten()

    # 获取测试集对应的日期索引
    test_dates = stock_data.index[-len(real_stock_price):]
    # 获取未来预测日期
    future_dates = pd.bdate_range(start=test_dates[-1] + pd.Timedelta(days=1), periods=forecast_days)

    # 绘图
    plt.figure(figsize=(14, 7))
    plt.plot(stock_data.index, stock_data['close'], color='red', label='All Real Stock Price')
    plt.plot(test_dates, predicted_stock_price, color='blue', label='Predicted Stock Price')
    plt.plot(future_dates, future_predictions, color='green', label='Future Predictions')

    if len(future_predictions) > 0:
        # 最高、最低、最后一天
        max_idx = np.argmax(future_predictions)
        min_idx = np.argmin(future_predictions)
        max_date = future_dates[max_idx]
        min_date = future_dates[min_idx]
        max_price = future_predictions[max_idx]
        min_price = future_predictions[min_idx]
        last_date = future_dates[-1]
        last_price = future_predictions[-1]

        plt.annotate(f"最高: {max_price:.2f}",
                     xy=(max_date, max_price),
                     xytext=(0, 20),
                     textcoords="offset points",
                     xycoords='data',
                     fontsize=12, color='green',
                     bbox=dict(boxstyle="round,pad=0.2", fc="yellow", alpha=0.5))
        plt.annotate(f"最低: {min_price:.2f}",
                     xy=(min_date, min_price),
                     xytext=(0, -30),
                     textcoords="offset points",
                     xycoords='data',
                     fontsize=12, color='green',
                     bbox=dict(boxstyle="round,pad=0.2", fc="yellow", alpha=0.5))
        plt.annotate(f"{last_price:.2f}", 
                     xy=(last_date, last_price), 
                     xytext=(0, 10),
                     textcoords="offset points",
                     xycoords='data',
                     fontsize=12, color='green',
                     bbox=dict(boxstyle="round,pad=0.2", fc="yellow", alpha=0.5))

    plt.title(f'{stock_code} Stock Price Prediction ({time_range})')
    plt.xlabel('Date')
    plt.ylabel('Stock Price')
    plt.legend()
    plt.tight_layout()

    img_io = io.BytesIO()
    plt.savefig(img_io, format='PNG')
    plt.close()
    img_io.seek(0)
    img = Image.open(img_io)

    return mse, img

# Gradio 界面函数
def gradio_interface(stock_code, time_range, custom_code, custom_time):
    code = custom_code if stock_code == "自定义输入" else stock_code
    # 判断时间范围
    if time_range == "自定义输入" and custom_time:
        try:
            start_date, end_date = [s.strip() for s in custom_time.split(",")]
        except Exception:
            return "自定义时间格式错误，应为: 20200101,20231231", None
    else:
        start_date, end_date = None, None
    mse, img = train_and_predict(code, time_range, start_date, end_date)
    mse_text = f'均方误差 (Mean Squared Error): {mse:.4f}'
    return mse_text, img

def show_custom_input(stock_code):
    return gr.update(visible=(stock_code == "自定义输入"))

def show_custom_time(time_range):
    return gr.update(visible=(time_range == "自定义输入"))

fixed_stocks = ['002027.SZ', '601288.SH', '601318.SH', '600900.SH', '601398.SH']

with gr.Blocks() as interface:
    with gr.Row():
        with gr.Column(scale=1):
            stock_dropdown = gr.Dropdown(
                choices=fixed_stocks + ['自定义输入'],
                label="选择股票",
                value='002027.SZ'
            )
            custom_code = gr.Textbox(label="输入股票代码", visible=False)
            time_range = gr.Dropdown(
                choices=['一年', '三年', '五年', '十年', '全部', '自定义输入'],
                label="时间范围",
                value='全部'
            )
            custom_time = gr.Textbox(label="输入起止日期（格式: 20200101,20231231）", visible=False)
            predict_btn = gr.Button("预测")
        with gr.Column(scale=2):
            output_text = gr.Textbox(label="均方误差")
            output_img = gr.Image(label="预测图")

    stock_dropdown.change(show_custom_input, stock_dropdown, custom_code)
    time_range.change(show_custom_time, time_range, custom_time)
    predict_btn.click(
        gradio_interface,
        inputs=[stock_dropdown, time_range, custom_code, custom_time],
        outputs=[output_text, output_img]
    )

interface.launch(share=True)